import time
import torch
import numpy as np
from importlib import import_module
import argparse
import utils
import train
parser = argparse.ArgumentParser(description='Bruce-Bert-Text-Classification')
parser.add_argument('--model',type=str,default='BruceERNIEDPCNN',help='choose a model BruceBert, BruceBertCNN, '
                                                                 'BruceBertRNN, BruceBertDPCNN, BruceERNIEDPCNN')
args = parser.parse_args()


if __name__ == '__main__':
    dataset = 'THUCNews'
    model_name = args.model # 这里args.model的model与上述 --model 这里的model名字对应； 得到的model_name是上述default被赋予的名字
    # print(model_name)
    x = import_module('models.' + model_name)
    config = x.Config(dataset)
    # start 保证每次运行结果一样
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed(1)
    torch.backends.cudnn.deterministic = True
    # end 保证每次运行结果一样

    start_time = time.time()
    print("加载数据")
    '''No.1 去写build_dataset'''
    train_data,dev_data,test_data = utils.build_dataset(config)
    '''No.3 去写build_iterator'''
    train_iter = utils.build_iterator(train_data, config)
    # for i, (trains, labels) in enumerate(train_iter):
    #     print(i, labels)
    dev_iter = utils.build_iterator(dev_data, config)
    test_iter = utils.build_iterator(test_data, config)

    time_dif = utils.get_time_dif(start_time)
    print("模型开始之前，准备数据时间: ", time_dif)

    # 模型训练、评估、测试
    model = x.Model(config).to(config.device)
    train.train(config, model, train_iter, dev_iter, test_iter)
    # train.test(config, model, test_iter)